# Created on 2018/12
# Author: Kaituo XU
import argparse

import torch
import torch.nn  as nn
import torch.nn.functional as F
import numpy as np
from utils import overlap_and_add
from mindspore import Tensor
import numpy as np
import mindspore
import torch

x = torch.ones(1, 2, 5, 1)
t = torch.ones(1, 2, 10, 1)
#print(t)
index = torch.tensor([0, 1, 1, 2, 2, 3, 3, 4, 4, 4])
#print(index)
x.index_add_(-2, index, t)
print(x)
print("*" * 7)
index = torch.tensor([0, 1, 1, 2, 2, 3, 3, 4, 4, 4])
#print(index)
index1 = torch.unsqueeze(index, 1)
index1 = torch.unsqueeze(index1, 0)
#print(index1)
index2 = index1
index3 = torch.cat((index1, index2),0)
#print(index3)
index3 = torch.unsqueeze(index3, 0)
# print(index3)
# print(index3.shape)

# index2 = torch.unsqueeze(index, 1)
# print(index2)
# index2 = index2.repeat(1, 3)
# print(index2)
y = torch.ones(1, 2, 5, 1)
y.scatter_add_(2, index3, t)
print(y)


# index2 = torch.unsqueeze(index, 1)
# print(index2)
# index2 = index2.repeat(1, 3)
# print(index2)
# y = torch.ones(5, 3)
# y.scatter_add_(0, index2, t)
# print(y)
#
# gamma = Tensor(np.ones([1, 32, 1]), mindspore.float32)
# print(gamma)
# gamma = mindspore.Parameter(mindspore.Tensor((1, 32, 1)))
# print(gamma.shape)
# gamma1 = nn.Parameter(torch.Tensor(1, 32, 1))
# t = gamma1.size()
# print(t)
# a = mindspore.Tensor((1, 32, 1))
# print(a.shape)
# b = torch.Tensor(1, 32, 1)
# print(b.shape)
# x = np.arange(12).reshape(3,4)
# y = Tensor(x)
# z = y.shape
# print(z)
# def convolve_overlap_add(a, h):
#     a_length = len(a)
#     h_length = len(h)
#     o_length = h_length-1
#     # blocks = get_blocks(a, h_length)
#     # length = blocks.shape[1]
#     # hm = generate_convolve_matrix(h, length)
#     # cb = convolve_blocks(blocks, hm)
#
#     result = np.zeros(length * length)
#     read_offset = 0
#     overlap_add = np.zeros(length)
#     for i in range(cb.shape[0]):
#         if i == 0:
#             result[0:length] = cb[i]
#             read_offset = h_length
#         else:
#             result[read_offset:read_offset + length] = overlap_add + cb[i]
#             read_offset = read_offset + h_length
#         overlap_add = cb[i][h_length:]
#         overlap_add = np.pad(overlap_add, (0, length - len(overlap_add)), 'constant', constant_values=(0, 0))
#     return result[:a_length+o_length]

#
# def convolve_overlap_add(a,h)
#     a_length = len(a)
#     h_length = len(h)
#     o_length = h_length-1
#     blocks = get_blocks(a, h_length)
#     length = blocks.shape[1]
#     hm = generate_convolve_matrix(h, length)
#     cb = convolve_blocks(blocks, hm)
#
#     result = np.zeros(length *length)
#     read_offset = 0
#     overlap_add = np.zeros(length)
#     for i in range(cb.shape[0]):
#          if i == 0:
#              result[0:length] = cb[i]
#              read_offset = h_length
#          else:
#              result[read_offset:read_offset + length] = overlap_add + cb[i]
#              read_offset = read_offset + h_length
#          overlap_add = cb[i][h_length:]
#
#     return result[:a_length+o_length]
#
#
# gamma = nn.Parameter(torch.Tensor(1, 32, 1))
# beta = nn.Parameter(torch.Tensor(1, 32, 1))
# mean = y.mean(dim=1, keepdim=True).mean(dim=2, keepdim=True)
# var = (torch.pow(y-mean, 2)).mean(dim=1, keepdim=True).mean(dim=2, keepdim=True)
x = torch.ones(1, 2, 5, 1)
t = torch.ones(1, 2, 10, 1)
index = torch.tensor([0, 1, 1, 2, 2, 3, 3, 4, 4, 4])
x.index_add_(-2, index, t)
print(x)
print("*" * 7)
index = torch.tensor([0, 1, 1, 2, 2, 3, 3, 4, 4, 4])
index1 = torch.unsqueeze(index, 1)
index1 = torch.unsqueeze(index1, 0)
index2 = index1
index3 = torch.cat((index1, index2),0)
index3 = torch.unsqueeze(index3, 0)
y = torch.ones(1, 2, 5, 1)
y.scatter_add_(2, index3, t)
print(y)


